# 此函数用于对称分布傅里叶的非线性部分的仿真，NL + 微扰
# 传入参数：
# gamma: 非线性系数，1 / (W * km)
# polar_n: 偏振数
# deltaz: 每个span进行分布傅里叶的长度
from filecmp import cmp

import numpy as np
from numpy import real, imag


def nonlinearity_step(signal, Signal, Fiber_para):
    data = signal
    # ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** *信号参数 ** ** ** ** ** ** ** ** ** ** ** ** ** ** **
    polar_n = Signal["polarization"]
    # ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** *光纤参数 ** ** ** ** ** ** ** ** ** ** ** ** ** ** ** **
    deltaz = Fiber_para["step_size"]  # 步长[km]
    gamma = Fiber_para["gamma"]  # 非线性系数[1/(W*km)]
    steptype = Fiber_para['ssfm_type']
    delta_z = Fiber_para['step_size']
    if polar_n == 1:  # 单偏振

        data = data * np.exp(1j * gamma * (abs(data) ** 2) * delta_z)
        sig = data
    # == == == == == == == == == == == == == == == == == == == Manakov == == == == == == == == == == == == == == == == =
    if polar_n == 2:  # 双偏振

        # 光纤非线性
        data_x = data[0:, 0]
        data_y = data[0:, 1]
        # X 偏振
        data_x = data_x * np.exp(8 / 9 * 1j * gamma * ((abs(data_y) ** 2) + (abs(data_x) ** 2)) * delta_z)
        # Y 偏振
        data_y = data_y * np.exp(8 / 9 * 1j * gamma * ((abs(data_y) ** 2) + (abs(data_x) ** 2)) * delta_z)

        data_x = data_x.reshape(-1, 1)
        data_y = data_y.reshape(-1, 1)

        sig = np.concatenate((data_x, data_y), axis=1)

        # print('非线性后功率')
        # print(np.mean(np.abs(sig[0:, 0]) ** 2))
    return sig
